import tushare as ts
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 初始化tushare接口
pro = ts.pro_api('137d8709cb300cf567f3bc419e6ef6d196abef75f889ed1355dc9896')

# 拉取数据
df = pro.daily(**{
    "ts_code": "000002.SZ",
    "trade_date": "",
    "start_date": "",
    "end_date": "",
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code", "trade_date", "open", "high", "low", "close",
    "pre_close", "change", "pct_chg", "vol", "amount"
])


# 数据预处理
def preprocess_data(df):
    # 转换日期格式并排序
    df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')
    df = df.sort_values('trade_date').reset_index(drop=True)

    # 计算技术指标
    # 移动平均线
    df['SMA_5'] = df['close'].rolling(5).mean()
    df['SMA_10'] = df['close'].rolling(10).mean()
    df['SMA_20'] = df['close'].rolling(20).mean()

    # RSI
    delta = df['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.rolling(14, min_periods=1).mean()
    avg_loss = loss.rolling(14, min_periods=1).mean()
    rs = avg_gain / avg_loss
    df['RSI'] = 100 - (100 / (1 + rs))

    # MACD
    df['EMA12'] = df['close'].ewm(span=12, adjust=False).mean()
    df['EMA26'] = df['close'].ewm(span=26, adjust=False).mean()
    df['MACD'] = df['EMA12'] - df['EMA26']
    df['Signal_Line'] = df['MACD'].ewm(span=9, adjust=False).mean()
    df['MACD_Hist'] = df['MACD'] - df['Signal_Line']

    # 布林带
    df['BB_SMA_20'] = df['close'].rolling(20).mean()
    df['BB_std_20'] = df['close'].rolling(20).std()
    df['BB_Upper'] = df['BB_SMA_20'] + 2 * df['BB_std_20']
    df['BB_Lower'] = df['BB_SMA_20'] - 2 * df['BB_std_20']

    # 成交量变化率
    df['Volume_Change'] = df['vol'].pct_change()

    # 删除缺失值
    df = df.dropna()

    # 目标变量（下一日涨跌）
    df['target'] = (df['close'].shift(-1) > df['close']).astype(int)
    df = df.dropna()  # 删除最后一行

    return df


# 特征工程
df = preprocess_data(df)
features = ['SMA_5', 'SMA_10', 'RSI', 'MACD_Hist', 'Volume_Change']
X = df[features]
y = df['target']

# 数据分割（修复索引对齐问题的核心部分）
split_idx = int(len(X) * 0.8)
train_end = split_idx
test_start = split_idx

# 训练集：前80%的数据
X_train = X.iloc[:train_end]
y_train = y.iloc[:train_end]

# 测试集：后20%的数据（排除最后一行）
X_test = X.iloc[test_start:-1]  
y_test = y.iloc[test_start:-1]

# 模型训练
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 模型预测
y_pred = model.predict(X_test)

# 策略计算（对齐的日期和收益率）
test_dates = df['trade_date'].iloc[test_start + 1:]  # 正确对齐
test_returns = df['pct_chg'].iloc[test_start + 1:].reset_index(drop=True)

# 策略收益率
strategy_returns = test_returns.where(y_pred == 1, 0)
cumulative_strategy = (1 + strategy_returns / 100).cumprod().iloc[-1] - 1

# 基准收益率
benchmark_return = (1 + test_returns / 100).cumprod().iloc[-1] - 1

# 评估结果
print(f"模型准确率：{accuracy_score(y_test, y_pred):.2%}")
print(f"策略累计收益率：{cumulative_strategy:.2%}")
print(f"基准累计收益率：{benchmark_return:.2%}")


# 可视化部分
def plot_results():
    # 特征重要性
    plt.figure(figsize=(10, 6))
    pd.Series(model.feature_importances_, index=features).sort_values().plot(kind='barh')
    plt.title('特征重要性')

    # 收益对比
    plt.figure(figsize=(12, 6))
    plt.plot(test_dates, (1 + test_returns / 100).cumprod(), label='基准')
    plt.plot(test_dates, (1 + strategy_returns / 100).cumprod(), label='策略')
    plt.legend()
    plt.title('收益对比')

    # 预测对比
    plt.figure(figsize=(12, 6))
    plt.plot(test_dates, y_test.reset_index(drop=True), label='实际值', color='blue', alpha=0.5)
    plt.plot(test_dates, y_pred, label='预测值', color='red', alpha=0.5)
    plt.legend()
    plt.title('预测对比')


plot_results()
plt.show()